package com.spring.distribute.zk.config;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.connection.StandardConnectionHandlingPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.state.SessionConnectionStateErrorPolicy;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import java.util.concurrent.TimeUnit;

@Slf4j
@Configuration
public class ZkConfig {

    @Value("${zk.address}")
    private String zkAddress;

    @Value("${zk.connectTimeoutMs}")
    public void setConnectTimeoutMs(Integer ms) {
        connectTimeoutMs  = ms;
    }
    private static Integer connectTimeoutMs;

    @Value("${zk.sessionTimeoutMs}")
    private Integer sessionTimeoutMs;

    @Value("${zk.maxRetry}")
    private Integer maxRetry;

    @Bean(name = "distributeCuratorFramework")
    public CuratorFramework rootCuratorFramework() {
        CuratorFramework curatorFramework = CuratorFrameworkFactory.builder()
            .connectString(this.zkAddress)
            .sessionTimeoutMs(this.sessionTimeoutMs)
            .connectionTimeoutMs(this.connectTimeoutMs)
            .connectionHandlingPolicy(new StandardConnectionHandlingPolicy(50))
            .connectionStateErrorPolicy(new SessionConnectionStateErrorPolicy())
            .retryPolicy(new ExponentialBackoffRetry(1000, this.maxRetry))
            .build();

        try {
            curatorFramework.start();
            int loop = 0;
            do {
                TimeUnit.MILLISECONDS.sleep(500);
            } while (!curatorFramework.blockUntilConnected(this.connectTimeoutMs, TimeUnit.MILLISECONDS) && loop++ < 5);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return curatorFramework;
    }
}
